{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "terminal-example",
   "metadata": {},
   "source": [
    "# Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "patent-lightning",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.autograd as autograd         # computation graph\n",
    "from torch import Tensor                  # tensor node in the computation graph\n",
    "import torch.nn as nn                     # neural networks\n",
    "import torch.optim as optim               # optimizers e.g. gradient descent, ADAM, etc.\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import matplotlib.ticker\n",
    "\n",
    "import numpy as np\n",
    "import time\n",
    "from pyDOE import lhs         #Latin Hypercube Sampling\n",
    "import scipy.io\n",
    "\n",
    "#Set default dtype to float32\n",
    "torch.set_default_dtype(torch.float)\n",
    "\n",
    "#PyTorch random number generator\n",
    "torch.manual_seed(1234)\n",
    "\n",
    "# Random number generators in other libraries\n",
    "np.random.seed(1234)\n",
    "\n",
    "# Device configuration\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "print(device)\n",
    "\n",
    "if device == 'cuda': \n",
    "    print(torch.cuda.get_device_name()) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "documentary-jewelry",
   "metadata": {},
   "source": [
    "# *Data Prep*\n",
    "\n",
    "Training and Testing data is prepared from the solution file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "peripheral-joseph",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = scipy.io.loadmat('Data/burgers_shock_mu_01_pi.mat')  \t# Load data from file\n",
    "x = data['x']                                   # 256 points between -1 and 1 [256x1]\n",
    "t = data['t']                                   # 100 time points between 0 and 1 [100x1] \n",
    "usol = data['usol']                             # solution of 256x100 grid points\n",
    "\n",
    "X, T = np.meshgrid(x,t)                         # makes 2 arrays X and T such that u(X[i],T[j])=usol[i][j] are a tuple"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "asian-driver",
   "metadata": {},
   "source": [
    "# Test Data\n",
    "\n",
    "We prepare the test data to compare against the solution produced by the PINN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "interesting-clerk",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yddknKA2Xohp"
   },
   "outputs": [],
   "source": [
    "''' X_u_test = [X[i],T[i]] [25600,2] for interpolation'''\n",
    "X_u_test = np.hstack((X.flatten()[:,None], T.flatten()[:,None]))\n",
    "\n",
    "# Domain bounds\n",
    "lb = X_u_test[0]  # [-1. 0.]\n",
    "ub = X_u_test[-1] # [1.  0.99]\n",
    "\n",
    "'''\n",
    "   Fortran Style ('F') flatten,stacked column wise!\n",
    "   u = [c1 \n",
    "        c2\n",
    "        .\n",
    "        .\n",
    "        cn]\n",
    "\n",
    "   u =  [25600x1] \n",
    "'''\n",
    "u_true = usol.flatten('F')[:,None] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "middle-audit",
   "metadata": {},
   "source": [
    "# Training Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "stainless-thesaurus",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8UVJmvZbXjXb"
   },
   "outputs": [],
   "source": [
    "def trainingdata(N_u,N_f):\n",
    "\n",
    "    '''Boundary Conditions'''\n",
    "\n",
    "    #Initial Condition -1 =< x =<1 and t = 0  \n",
    "    leftedge_x = np.hstack((X[0,:][:,None], T[0,:][:,None])) #L1\n",
    "    leftedge_u = usol[:,0][:,None]\n",
    "\n",
    "    #Boundary Condition x = -1 and 0 =< t =<1\n",
    "    bottomedge_x = np.hstack((X[:,0][:,None], T[:,0][:,None])) #L2\n",
    "    bottomedge_u = usol[-1,:][:,None]\n",
    "\n",
    "    #Boundary Condition x = 1 and 0 =< t =<1\n",
    "    topedge_x = np.hstack((X[:,-1][:,None], T[:,0][:,None])) #L3\n",
    "    topedge_u = usol[0,:][:,None]\n",
    "\n",
    "    all_X_u_train = np.vstack([leftedge_x, bottomedge_x, topedge_x]) # X_u_train [456,2] (456 = 256(L1)+100(L2)+100(L3))\n",
    "    all_u_train = np.vstack([leftedge_u, bottomedge_u, topedge_u])   #corresponding u [456x1]\n",
    "\n",
    "    #choose random N_u points for training\n",
    "    idx = np.random.choice(all_X_u_train.shape[0], N_u, replace=False) \n",
    "\n",
    "    X_u_train = all_X_u_train[idx, :] #choose indices from  set 'idx' (x,t)\n",
    "    u_train = all_u_train[idx,:]      #choose corresponding u\n",
    "\n",
    "    '''Collocation Points'''\n",
    "\n",
    "    # Latin Hypercube sampling for collocation points \n",
    "    # N_f sets of tuples(x,t)\n",
    "    X_f_train = lb + (ub-lb)*lhs(2,N_f) \n",
    "    X_f_train = np.vstack((X_f_train, X_u_train)) # append training points to collocation points \n",
    "\n",
    "    return X_f_train, X_u_train, u_train \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "prescribed-spencer",
   "metadata": {},
   "source": [
    "# Physics Informed Neural Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "elect-austria",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sequentialmodel(nn.Module):\n",
    "    \n",
    "    def __init__(self,layers):\n",
    "        super().__init__() #call __init__ from parent class \n",
    "              \n",
    "        'activation function'\n",
    "        self.activation = nn.Tanh()\n",
    "\n",
    "        'loss function'\n",
    "        self.loss_function = nn.MSELoss(reduction ='mean')\n",
    "    \n",
    "        'Initialise neural network as a list using nn.Modulelist'  \n",
    "        self.linears = nn.ModuleList([nn.Linear(layers[i], layers[i+1]) for i in range(len(layers)-1)])\n",
    "        \n",
    "        self.iter = 0\n",
    "        \n",
    "        '''\n",
    "        Alternatively:\n",
    "        \n",
    "        *all layers are callable \n",
    "    \n",
    "        Simple linear Layers\n",
    "        self.fc1 = nn.Linear(2,50)\n",
    "        self.fc2 = nn.Linear(50,50)\n",
    "        self.fc3 = nn.Linear(50,50)\n",
    "        self.fc4 = nn.Linear(50,1)\n",
    "        \n",
    "        '''\n",
    "    \n",
    "        'Xavier Normal Initialization'\n",
    "        # std = gain * sqrt(2/(input_dim+output_dim))\n",
    "        for i in range(len(layers)-1):\n",
    "            \n",
    "            # weights from a normal distribution with \n",
    "            # Recommended gain value for tanh = 5/3?\n",
    "            nn.init.xavier_normal_(self.linears[i].weight.data, gain=1.0)\n",
    "            \n",
    "            # set biases to zero\n",
    "            nn.init.zeros_(self.linears[i].bias.data)\n",
    "            \n",
    "    'foward pass'\n",
    "    def forward(self,x):\n",
    "        \n",
    "        if torch.is_tensor(x) != True:         \n",
    "            x = torch.from_numpy(x)                \n",
    "        \n",
    "        u_b = torch.from_numpy(ub).float().to(device)\n",
    "        l_b = torch.from_numpy(lb).float().to(device)\n",
    "                      \n",
    "        #preprocessing input \n",
    "        x = (x - l_b)/(u_b - l_b) #feature scaling\n",
    "        \n",
    "        #convert to float\n",
    "        a = x.float()\n",
    "                        \n",
    "        '''     \n",
    "        Alternatively:\n",
    "        \n",
    "        a = self.activation(self.fc1(a))\n",
    "        a = self.activation(self.fc2(a))\n",
    "        a = self.activation(self.fc3(a))\n",
    "        a = self.fc4(a)\n",
    "        \n",
    "        '''\n",
    "        \n",
    "        for i in range(len(layers)-2):\n",
    "            \n",
    "            z = self.linears[i](a)\n",
    "                        \n",
    "            a = self.activation(z)\n",
    "            \n",
    "        a = self.linears[-1](a)\n",
    "        \n",
    "        return a\n",
    "                        \n",
    "    def loss_BC(self,x,y):\n",
    "                \n",
    "        loss_u = self.loss_function(self.forward(x), y)\n",
    "                \n",
    "        return loss_u\n",
    "    \n",
    "    def loss_PDE(self, x_to_train_f):\n",
    "        \n",
    "        nu = 0.01/np.pi\n",
    "                \n",
    "        x_1_f = x_to_train_f[:,[0]]\n",
    "        x_2_f = x_to_train_f[:,[1]]\n",
    "                        \n",
    "        g = x_to_train_f.clone()\n",
    "                        \n",
    "        g.requires_grad = True\n",
    "        \n",
    "        u = self.forward(g)\n",
    "                \n",
    "        u_x_t = autograd.grad(u,g,torch.ones([x_to_train_f.shape[0], 1]).to(device), retain_graph=True, create_graph=True)[0]\n",
    "                                \n",
    "        u_xx_tt = autograd.grad(u_x_t,g,torch.ones(x_to_train_f.shape).to(device), create_graph=True)[0]\n",
    "                                                            \n",
    "        u_x = u_x_t[:,[0]]\n",
    "        \n",
    "        u_t = u_x_t[:,[1]]\n",
    "        \n",
    "        u_xx = u_xx_tt[:,[0]]\n",
    "                                        \n",
    "        f = u_t + (self.forward(g))*(u_x) - (nu)*u_xx \n",
    "        \n",
    "        loss_f = self.loss_function(f,f_hat)\n",
    "                \n",
    "        return loss_f\n",
    "    \n",
    "    def loss(self,x,y,x_to_train_f):\n",
    "\n",
    "        loss_u = self.loss_BC(x,y)\n",
    "        loss_f = self.loss_PDE(x_to_train_f)\n",
    "        \n",
    "        loss_val = loss_u + loss_f\n",
    "        \n",
    "        return loss_val\n",
    "     \n",
    "    'callable for optimizer'                                       \n",
    "    def closure(self):\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        loss = self.loss(X_u_train, u_train, X_f_train)\n",
    "        \n",
    "        loss.backward()\n",
    "                \n",
    "        self.iter += 1\n",
    "        \n",
    "        if self.iter % 100 == 0:\n",
    "\n",
    "            error_vec, _ = PINN.test()\n",
    "        \n",
    "            print(loss,error_vec)\n",
    "\n",
    "        return loss        \n",
    "    \n",
    "    'test neural network'\n",
    "    def test(self):\n",
    "                \n",
    "        u_pred = self.forward(X_u_test_tensor)\n",
    "        \n",
    "        error_vec = torch.linalg.norm((u-u_pred),2)/torch.linalg.norm(u,2)        # Relative L2 Norm of the error (Vector)\n",
    "        \n",
    "        u_pred = u_pred.cpu().detach().numpy()\n",
    "        \n",
    "        u_pred = np.reshape(u_pred,(256,100),order='F')\n",
    "                \n",
    "        return error_vec, u_pred"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "qualified-ceiling",
   "metadata": {
    "colab_type": "text",
    "id": "bOjuHdzAhib-"
   },
   "source": [
    "# *Solution Plot*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "entire-rolling",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UWqNuRMLhg4m"
   },
   "outputs": [],
   "source": [
    "def solutionplot(u_pred,X_u_train,u_train):\n",
    "    \n",
    "    fig, ax = plt.subplots()\n",
    "    ax.axis('off')\n",
    "\n",
    "    gs0 = gridspec.GridSpec(1, 2)\n",
    "    gs0.update(top=1-0.06, bottom=1-1/3, left=0.15, right=0.85, wspace=0)\n",
    "    ax = plt.subplot(gs0[:, :])\n",
    "\n",
    "    h = ax.imshow(u_pred, interpolation='nearest', cmap='rainbow', \n",
    "                extent=[T.min(), T.max(), X.min(), X.max()], \n",
    "                origin='lower', aspect='auto')\n",
    "    divider = make_axes_locatable(ax)\n",
    "    cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "    fig.colorbar(h, cax=cax)\n",
    "    \n",
    "    ax.plot(X_u_train[:,1], X_u_train[:,0], 'kx', label = 'Data (%d points)' % (u_train.shape[0]), markersize = 4, clip_on = False)\n",
    "\n",
    "    line = np.linspace(x.min(), x.max(), 2)[:,None]\n",
    "    ax.plot(t[25]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "    ax.plot(t[50]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "    ax.plot(t[75]*np.ones((2,1)), line, 'w-', linewidth = 1)    \n",
    "\n",
    "    ax.set_xlabel('$t$')\n",
    "    ax.set_ylabel('$x$')\n",
    "    ax.legend(frameon=False, loc = 'best')\n",
    "    ax.set_title('$u(x,t)$', fontsize = 10)\n",
    "    \n",
    "    ''' \n",
    "    Slices of the solution at points t = 0.25, t = 0.50 and t = 0.75\n",
    "    '''\n",
    "    \n",
    "    ####### Row 1: u(t,x) slices ##################\n",
    "    gs1 = gridspec.GridSpec(1, 3)\n",
    "    gs1.update(top=1-1/3, bottom=0, left=0.1, right=0.9, wspace=0.5)\n",
    "\n",
    "    ax = plt.subplot(gs1[0, 0])\n",
    "    ax.plot(x,usol.T[25,:], 'b-', linewidth = 2, label = 'Exact')       \n",
    "    ax.plot(x,u_pred.T[25,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "    ax.set_xlabel('$x$')\n",
    "    ax.set_ylabel('$u(x,t)$')    \n",
    "    ax.set_title('$t = 0.25s$', fontsize = 10)\n",
    "    ax.axis('square')\n",
    "    ax.set_xlim([-1.1,1.1])\n",
    "    ax.set_ylim([-1.1,1.1])\n",
    "\n",
    "    ax = plt.subplot(gs1[0, 1])\n",
    "    ax.plot(x,usol.T[50,:], 'b-', linewidth = 2, label = 'Exact')       \n",
    "    ax.plot(x,u_pred.T[50,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "    ax.set_xlabel('$x$')\n",
    "    ax.set_ylabel('$u(x,t)$')\n",
    "    ax.axis('square')\n",
    "    ax.set_xlim([-1.1,1.1])\n",
    "    ax.set_ylim([-1.1,1.1])\n",
    "    ax.set_title('$t = 0.50s$', fontsize = 10)\n",
    "    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.35), ncol=5, frameon=False)\n",
    "\n",
    "    ax = plt.subplot(gs1[0, 2])\n",
    "    ax.plot(x,usol.T[75,:], 'b-', linewidth = 2, label = 'Exact')       \n",
    "    ax.plot(x,u_pred.T[75,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "    ax.set_xlabel('$x$')\n",
    "    ax.set_ylabel('$u(x,t)$')\n",
    "    ax.axis('square')\n",
    "    ax.set_xlim([-1.1,1.1])\n",
    "    ax.set_ylim([-1.1,1.1])    \n",
    "    ax.set_title('$t = 0.75s$', fontsize = 10)\n",
    "    \n",
    "    plt.savefig('Burgers.png',dpi = 500)   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "broad-house",
   "metadata": {},
   "source": [
    "# Main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "private-still",
   "metadata": {},
   "outputs": [],
   "source": [
    "'Generate Training data'\n",
    "N_u = 100 #Total number of data points for 'u'\n",
    "N_f = 10000 #Total number of collocation points \n",
    "X_f_train_np_array, X_u_train_np_array, u_train_np_array = trainingdata(N_u,N_f)\n",
    "\n",
    "'Convert to tensor and send to GPU'\n",
    "X_f_train = torch.from_numpy(X_f_train_np_array).float().to(device)\n",
    "X_u_train = torch.from_numpy(X_u_train_np_array).float().to(device)\n",
    "u_train = torch.from_numpy(u_train_np_array).float().to(device)\n",
    "X_u_test_tensor = torch.from_numpy(X_u_test).float().to(device)\n",
    "u = torch.from_numpy(u_true).float().to(device)\n",
    "f_hat = torch.zeros(X_f_train.shape[0],1).to(device)\n",
    "\n",
    "layers = np.array([2,20,20,20,20,20,20,20,20,1]) #8 hidden layers\n",
    "\n",
    "PINN = Sequentialmodel(layers)\n",
    "       \n",
    "PINN.to(device)\n",
    "\n",
    "'Neural Network Summary'\n",
    "print(PINN)\n",
    "\n",
    "params = list(PINN.parameters())\n",
    "\n",
    "'''Optimization'''\n",
    "\n",
    "'L-BFGS Optimizer'\n",
    "optimizer = torch.optim.LBFGS(PINN.parameters(), lr=0.1, \n",
    "                              max_iter = 250, \n",
    "                              max_eval = None, \n",
    "                              tolerance_grad = 1e-05, \n",
    "                              tolerance_change = 1e-09, \n",
    "                              history_size = 100, \n",
    "                              line_search_fn = 'strong_wolfe')\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "optimizer.step(PINN.closure)\n",
    "\n",
    "\n",
    "'Adam Optimizer'\n",
    "# optimizer = optim.Adam(PINN.parameters(), lr=0.001,betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)\n",
    "\n",
    "# max_iter = 20000\n",
    "\n",
    "# start_time = time.time()\n",
    "\n",
    "# for i in range(max_iter):\n",
    "\n",
    "#     loss = PINN.loss(X_u_train, u_train, X_f_train)\n",
    "           \n",
    "#     optimizer.zero_grad()     # zeroes the gradient buffers of all parameters\n",
    "    \n",
    "#     loss.backward() #backprop\n",
    "\n",
    "#     optimizer.step()\n",
    "    \n",
    "#     if i % (max_iter/10) == 0:\n",
    "\n",
    "#         error_vec, _ = PINN.test()\n",
    "\n",
    "#         print(loss,error_vec)\n",
    "    \n",
    "    \n",
    "elapsed = time.time() - start_time                \n",
    "print('Training time: %.2f' % (elapsed))\n",
    "\n",
    "\n",
    "''' Model Accuracy ''' \n",
    "error_vec, u_pred = PINN.test()\n",
    "\n",
    "print('Test Error: %.5f'  % (error_vec))\n",
    "\n",
    "\n",
    "''' Solution Plot '''\n",
    "# solutionplot(u_pred,X_u_train.cpu().detach().numpy(),u_train.cpu().detach().numpy())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
